#include "h264_file_nalu_pars.h"

static bool MatchStartCode3Bytes(unsigned char *Buf)
{
    if (Buf[0] != 0 || Buf[1] != 0 || Buf[2] != 1)
        return false; // 0x000001?
    else
        return true;
}
 
static bool MatchStartCode4Bytes(unsigned char *Buf)
{
    if (Buf[0] != 0 || Buf[1] != 0 || Buf[2] != 0 || Buf[3] != 1)
        return false; // 0x00000001?
    else
        return true;
}


int H264FileNaluPars::init(std::string path, fmt_type_t fmt_type)
{
    m_fmt_type = fmt_type;
    m_h264_fd = fopen(path.c_str(), "rb");
    if (m_h264_fd == nullptr) {
        perror("fopen");
        deinit();
        exit(EXIT_FAILURE);
    }

    m_nal_num = 0;
    m_data_offset = 0;

    printf("-----+---------------- NALU Table ------+---------+\n");
    printf(" NUM |    POS    |   SCL  |    IDC |  TYPE |   LEN   |\n");
    printf("-----+-----------+--------+--------+-------+---------+\n");

    return 0;
}

int H264FileNaluPars::deinit()
{
    if (m_h264_fd != nullptr) {
        fclose(m_h264_fd);
    }
    return 0;
}

int H264FileNaluPars::get_annexb_nalu()
{
    size_t data_buf_len = 10240;  // 以太网最大传输单元(MTU)常见的值是1500字节
    uint8_t data_buf[10240];
    
    // 从文件中读取一定量数据
    size_t read_len = fread(data_buf, 1, data_buf_len, m_h264_fd);
    if (read_len < 3) {
        // delete[] data_buf;
        return 0;
    }

    uint8_t *pdata = data_buf;
    unsigned char *pdata_end = data_buf + read_len;
 
    // 判断NAL帧头是否正确
    if (!MatchStartCode3Bytes(pdata)) {
        if (read_len < 4) {
            // delete[] data_buf;
            return 0;
        }
        if (!MatchStartCode4Bytes(pdata)) {
            // delete[] data_buf;
            return -1;
        } else {
            m_nalu_frame.statcode_length = 4;
            pdata += 4;
        }
    } else {
        m_nalu_frame.statcode_length = 3;
        pdata += 3;
    }

    int next_startcode_length = 0;
    ssize_t rewind = 0;
 
    while (pdata != pdata_end) {
        // 不为NALU Header则数据位向后移动 否则跳出循环
        if (!MatchStartCode4Bytes(pdata)) {
            if (MatchStartCode3Bytes(pdata)) {
                next_startcode_length = 3;
                break;
            }
        } else {
            next_startcode_length = 4;
            break;
        }
        ++pdata;
        
        // 如果缓冲中没有找到帧头，再读一个缓冲区的数据
        if (pdata == pdata_end) {
            if (feof(m_h264_fd)) {
                // 文件读取完毕
                break;
            }
 
            rewind = -3;
            if (0 != fseek(m_h264_fd, rewind, SEEK_CUR)) {
                // delete[] data_buf;
                printf("GetAnnexbNALU: Cannot fseek in the bit stream file");
                return -1;
            }
 
            pdata -= 3;
 
            ssize_t cur_data_len = pdata - data_buf;
            m_nalu_frame.buf.insert(m_nalu_frame.buf.end(), std::begin(data_buf), std::begin(data_buf)+cur_data_len);
 
            read_len = fread(data_buf, 1, data_buf_len, m_h264_fd);
            if (read_len <= 0) {
                // delete[] data_buf;
                return -1;
            }
            pdata_end = data_buf + read_len; // update
            pdata = data_buf;
        }
    }
 
    // 缓冲中找到了另外一个开始头，
    rewind = pdata - pdata_end;  // 计算无效数据的长度
 
    if (0 != fseek(m_h264_fd, rewind, SEEK_CUR)) {
        // delete[] data_buf;
        printf("GetAnnexbNALU: Cannot fseek in the bit stream file");
        return -1;
    }
    
    // 更新一帧NALU数据
    // ssize_t cur_data_len = pdata - data_buf;  // 计算有效数据的长度
    // m_nalu_frame.buf.insert(m_nalu_frame.buf.end(), std::begin(data_buf), std::end(data_buf + cur_data_len));
    m_nalu_frame.buf.insert(m_nalu_frame.buf.end(), std::begin(data_buf), std::begin(data_buf) + (pdata - data_buf));

    m_nalu_frame.forbidden_bit = m_nalu_frame.buf[m_nalu_frame.statcode_length] & 0x80;     // 1 bit
    m_nalu_frame.nal_reference_idc = m_nalu_frame.buf[m_nalu_frame.statcode_length] & 0x60; // 2 bit

    if (FMT_H264 == m_fmt_type) {
        m_nalu_frame.nal_unit_type = (m_nalu_frame.buf[m_nalu_frame.statcode_length]) & 0x1f;   // 5 bit h264
    } else if (FMT_H265 == m_fmt_type) {
        m_nalu_frame.nal_unit_type = ((m_nalu_frame.buf[m_nalu_frame.statcode_length]) & 0x7E) >> 1;   // 5 bit h265
    }
    
    // 打印前10个字符
    // for (int i = 0; i < 10; ++i) {
    //     if (i >= m_nalu_frame.buf.size()) break;
    //     printf("%3x", m_nalu_frame.buf[i]);
    // }

    return m_nalu_frame.buf.size();
}

std::vector<uint8_t>& H264FileNaluPars::get_nalu_frame()
{
    int data_lenth;
    char type_str[20] = {0};
    char idc_str[20] = {0};

    m_nalu_frame.buf.clear();  // 清空上一帧缓冲

    if (!feof(m_h264_fd)) {
        
        data_lenth = get_annexb_nalu();
        if (data_lenth == 0)
            return m_nalu_frame.buf;
 
        switch (m_nalu_frame.nal_unit_type) {
        case NALU_TYPE_SLICE:
            sprintf(type_str, "SLICE");
            break;
        case NALU_TYPE_DPA:
            sprintf(type_str, "DPA");
            break;
        case NALU_TYPE_DPB:
            sprintf(type_str, "DPB");
            break;
        case NALU_TYPE_DPC:
            sprintf(type_str, "DPC");
            break;
        case NALU_TYPE_IDR:
            sprintf(type_str, "IDR");
            break;
        case NALU_TYPE_SEI:
            sprintf(type_str, "SEI");
            break;
        case NALU_TYPE_SPS:
            sprintf(type_str, "SPS");
            break;
        case NALU_TYPE_PPS:
            sprintf(type_str, "PPS");
            break;
        case NALU_TYPE_AUD:
            sprintf(type_str, "AUD");
            break;
        case NALU_TYPE_EOSEQ:
            sprintf(type_str, "EOSEQ");
            break;
        case NALU_TYPE_EOSTREAM:
            sprintf(type_str, "EOSTREAM");
            break;
        case NALU_TYPE_FILL:
            sprintf(type_str, "FILL");
            break;
        }

        switch (m_nalu_frame.nal_reference_idc >> 5) {
        case NALU_PRIORITY_DISPOSABLE:
            sprintf(idc_str, "DISPOS");
            break;
        case NALU_PRIRITY_LOW:
            sprintf(idc_str, "LOW");
            break;
        case NALU_PRIORITY_HIGH:
            sprintf(idc_str, "HIGH");
            break;
        case NALU_PRIORITY_HIGHEST:
            sprintf(idc_str, "HIGHEST");
            break;
        }

        fprintf(stdout, "%5d| 0x%08X| %7d| %7s| %6s| %8ld|\n", m_nal_num,
                m_data_offset, m_nalu_frame.statcode_length, idc_str, type_str, m_nalu_frame.buf.size());
 
        m_data_offset = m_data_offset + data_lenth;
        m_nal_num++;
 
        return m_nalu_frame.buf;
    }

    return m_nalu_frame.buf;
}
